use axum::{
    Router,
    response::sse::{Event, KeepAlive, Sse},
    routing::get,
};
use futures::{Stream, StreamExt};
use std::{convert::Infallible, time::Duration};
use tokio_stream::wrappers::IntervalStream;

#[tokio::main]
async fn main() {
    // 创建路由
    let app = Router::new().route("/sse", get(sse_handler));

    // 启动服务
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

// SSE 响应处理函数
async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    // 每秒发一个消息
    let stream = IntervalStream::new(tokio::time::interval(Duration::from_secs(1)))
        .enumerate()
        .map(|(i, _)| {
            let event = Event::default()
                .data(format!("这是第 {} 条消息", i))
                .event("message");
            Ok(event)
        });

    Sse::new(stream).keep_alive(KeepAlive::default())
}
